#!/usr/bin/env python
# -*- coding: UTF-8 -*-
import random
import matplotlib.pyplot as plt
from typing import Tuple, List

class TradingSimulator:
    def __init__(self, win_rate: float, risk_reward_ratio: float, initial_balance: float):
        self.win_rate = win_rate
        self.risk_reward_ratio = risk_reward_ratio
        self.balance = initial_balance
        self.base_risk_control_ratio = 0.005
        self.risk_control_level_list = [1, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 4, 4, 5, 5, 6, 7, 8, 9]
        self.risk_control_level = 1
        self.risk_control_ratio = self.base_risk_control_ratio
        self.trade_history = []
        self.balance_history = [self.balance]
        self.max_drawdown = 0
        self.max_drawdown_index = 0
        self.max_consecutive_wins = 0
        self.max_consecutive_wins_index = 0
        self.max_consecutive_losses = 0
        self.max_consecutive_losses_index = 0
        self.total_wins = 0
        self.total_losses = 0

    def simulate_trade(self, number_of_trades: int) -> Tuple[float, int, float, float, int, int, int, float]:
        high_water_mark = self.balance
        current_consecutive_wins = 0
        current_consecutive_losses = 0

        for i in range(number_of_trades):
            trade_result = random.random() < self.win_rate
            self.trade_history.append(trade_result)

            if trade_result:
                self.balance *= (1 + self.risk_control_ratio * self.risk_reward_ratio)
                self.risk_control_level = min(self.risk_control_level + 1, len(self.risk_control_level_list))
                self.risk_control_ratio = self.base_risk_control_ratio * self.risk_control_level_list[self.risk_control_level - 1]
                current_consecutive_wins += 1
                current_consecutive_losses = 0
                self.total_wins += 1
            else:
                self.balance *= (1 - self.risk_control_ratio)
                self.risk_control_level = max(1, self.risk_control_level - 1)
                self.risk_control_ratio = self.base_risk_control_ratio * self.risk_control_level_list[self.risk_control_level - 1]
                current_consecutive_losses += 1
                current_consecutive_wins = 0
                self.total_losses += 1

            self.balance_history.append(self.balance)
            drawdown = (high_water_mark - self.balance) / high_water_mark
            if drawdown > self.max_drawdown:
                self.max_drawdown = drawdown
                self.max_drawdown_index = i + 1

            self.max_consecutive_wins = max(self.max_consecutive_wins, current_consecutive_wins)
            self.max_consecutive_wins_index = max(self.max_consecutive_wins_index, i + 1) if self.max_consecutive_wins == current_consecutive_wins else self.max_consecutive_wins_index
            self.max_consecutive_losses = max(self.max_consecutive_losses, current_consecutive_losses)
            self.max_consecutive_losses_index = i + 1 if self.max_consecutive_losses == current_consecutive_losses else self.max_consecutive_losses_index

            high_water_mark = max(high_water_mark, self.balance)

        return (self.balance, self.risk_control_level, self.risk_control_ratio, self.max_drawdown,
                self.max_consecutive_wins, self.max_consecutive_losses, self.total_wins, self.total_losses)

    def plot_balance_curve(self, annotate: bool = False):
        plt.figure(figsize=(12, 6))
        plt.plot(self.balance_history, label='Account Balance', marker='o')
        if annotate:
            plt.scatter([self.max_drawdown_index], [self.balance_history[self.max_drawdown_index]], color='r', label='Max Drawdown')
            plt.annotate(f'Max Drawdown\n{self.balance_history[self.max_drawdown_index]:.2f}', (self.max_drawdown_index, self.balance_history[self.max_drawdown_index]), textcoords="offset points", xytext=(0,10), ha='center')
            plt.scatter([self.max_consecutive_wins_index], [self.balance_history[self.max_consecutive_wins_index]], color='g', label='Max Consecutive Wins')
            plt.annotate(f'Max Consecutive Wins\n{self.max_consecutive_wins}', (self.max_consecutive_wins_index, self.balance_history[self.max_consecutive_wins_index]), textcoords="offset points", xytext=(0,10), ha='center')
            plt.scatter([self.max_consecutive_losses_index], [self.balance_history[self.max_consecutive_losses_index]], color='b', label='Max Consecutive Losses')
            plt.annotate(f'Max Consecutive Losses\n{self.max_consecutive_losses}', (self.max_consecutive_losses_index, self.balance_history[self.max_consecutive_losses_index]), textcoords="offset points", xytext=(0,10), ha='center')
        plt.xlabel('Trade Number')
        plt.ylabel('Balance')
        plt.title('Account Balance Curve')
        plt.legend()
        plt.grid(True)
        plt.show()

# 使用示例
simulator = TradingSimulator(win_rate=0.5, risk_reward_ratio=1.5, initial_balance=100000)
final_balance, final_level, final_ratio, max_drawdown, max_wins, max_losses, total_wins, total_losses = simulator.simulate_trade(1000)

print(f"最终账户余额：{final_balance}")
print(f"最终风控等级：{final_level}")
print(f"最终风控比例：{final_ratio}")
print(f"最大回撤：{max_drawdown*100:.2f}%")
print(f"最大连续获胜次数：{max_wins}")
print(f"最大连续亏损次数：{max_losses}")
print(f"测试胜利次数：{total_wins}")
print(f"测试失败次数：{total_losses}")
print(f"测试胜率：{total_wins/(total_wins+total_losses)*100:.2f}%")

# 调用plot_balance_curve方法时，可以指定annotate参数来控制是否进行标注
simulator.plot_balance_curve(annotate=False)